package main.scala.algirithms.KDtree



import scala.collection.mutable

/**
 *
 * @param value 节点数据
 * @param dim   当前切分维度
 * @param left  左子结点
 * @param right 右子结点
 */
case class TreeNode(value: Seq[Double],
                    dim: Int,
                    var left: TreeNode,
                    var right: TreeNode) {

  var parent: TreeNode = _ //父结点
  var brotherNode: TreeNode = _ //兄弟结点

  if (left != null) {
    left.parent = this
    left.brotherNode = right
  }

  if (right != null) {
    right.parent = this
    right.brotherNode = left
  }

}





object KDtree {









  def main(args: Array[String]): Unit = {

    val nodes: Seq[Seq[Double]] =
      Seq(Seq(2, 3), Seq(5, 4), Seq(9, 6), Seq(4, 7), Seq(8, 1), Seq(7, 2))

    val treeNode: TreeNode = KdT.creatKdTree(nodes, 0, 2)

    println(treeNode)

    println(KdT.nearestSearch(treeNode, Seq(2.1, 4.5)).value)

    println("==============")
    nodes.map(x => {
      val d = KdT.euclidean(x, Seq(2.1, 4.5))
      (d, x)
    })
      .sortBy(_._1)
      .foreach(println)

    println("=======================================")


    KdT
      .knn(treeNode, Seq(2.1, 4.5), 3)
      .map(x => (x._1, x._2.value))
      .foreach(println)


    //    val path = "data\\classification\\kd_tree.txt";
//    val dataNodes = this.load(path, ",")
//    val dim = dataNodes(0).vec.length
//    val dimSet = mutable.HashSet[Int]()
//    0.until(dim).foreach(dimSet.add)
//    val kdTree = create(dataNodes, null, dimSet)
//    println(kdTree)
  }



}

object KdT{

  /**
   * 从root节点开始，DFS搜索直到叶子节点，同时在stack中顺序存储已经访问的节点。
   * 如果搜索到叶子节点，当前的叶子节点被设为最近邻节点。
   * 然后通过stack回溯:
   * 如果当前点的距离比最近邻点距离近，更新最近邻节点.
   * 然后检查以最近距离为半径的圆是否和父节点的超平面相交.
   * 如果相交，则必须到父节点的另外一侧，用同样的DFS搜索法，开始检查最近邻节点。
   * 如果不相交，则继续往上回溯，而父节点的另一侧子节点都被淘汰，不再考虑的范围中.
   * 当搜索回到root节点时，搜索完成，得到最近邻节点。
   *
   * @param treeNode
   * @param data
   * @param k
   * @return
   */
  def knn(treeNode: TreeNode, data: Seq[Double], k: Int): Array[(Double, TreeNode)] = {

    var resArr = new Array[(Double, TreeNode)](k)
      .map(_ => (Double.MaxValue, null))
      .asInstanceOf[Array[(Double, TreeNode)]]

    def finder(treeNode: TreeNode): TreeNode = {

      if (treeNode != null) {
        val dimr = data(treeNode.dim) - treeNode.value(treeNode.dim)
        if (dimr > 0) finder(treeNode.right) else finder(treeNode.left)

//        val distc: Double = distanceUtils.euclidean(treeNode.value, data)
        val distc: Double = euclidean(treeNode.value, data)

        if (distc < resArr.last._1  ) {
          resArr.update(k - 1, (distc, treeNode))
          resArr = resArr.sortBy(_._1)
        }

        if (math.abs(dimr) < resArr.last._1)
          if (dimr > 0) finder(treeNode.left) else finder(treeNode.right)

      }
      resArr.last._2
    }

    finder(treeNode)
    resArr

  }



  // 欧式距离算法
  def euclidean(p1: Seq[Double], p2: Seq[Double]) = {
    require(p1.size == p2.size)
    val d = p1
      .zip(p2)
      .map(tp => math.pow(tp._1 - tp._2, 2))
      .sum
    math.sqrt(d)
  }

  /**
   *
   * @param treeNode kdtree
   * @param data     查询点
   *                 最近邻搜索
   */
  def nearestSearch(treeNode: TreeNode, data: Seq[Double]): TreeNode = {

    var nearestNode: TreeNode = null //当前最近节点
    var minDist: Double = Double.MaxValue //当前最小距离

    def finder(treeNode: TreeNode): TreeNode = {

      treeNode match {
        case null => nearestNode
        case _ =>
          val dimr = data(treeNode.dim) - treeNode.value(treeNode.dim)
          if (dimr > 0) finder(treeNode.right) else finder(treeNode.left)

          val distc = euclidean(treeNode.value, data)
          if (distc <= minDist) {
            minDist = distc
            nearestNode = treeNode
          }

          // 目标点与当前节点相交
          if (math.abs(dimr) < minDist)
            if (dimr > 0) finder(treeNode.left) else finder(treeNode.right)

          nearestNode
      }
    }

    finder(treeNode)

  }

  /**
   *
   * @param value 数据序列
   * @param dim   当前划分的维度
   * @param shape 数据维数
   * @return
   */
  def creatKdTree(value: Seq[Seq[Double]], dim: Int, shape: Int): TreeNode = {

    // 数据按照当前划分的维度排序
//    val sorted = value.sortBy(_ (dim))
    val sorted = value.sortBy(x => x (dim))
    //中间位置的索引
    val midIndex: Int = value.length / 2

    sorted match {
      // 当节点为空时，返回null
      case Nil => null
      //节点不为空时，递归调用
      case _ =>
        val left = sorted.slice(0, midIndex)
        val right = sorted.slice(midIndex + 1, value.length)

        val leftNode = creatKdTree(left, (dim + 1) % shape, shape) //左子节点递归创建树
        val rightNode = creatKdTree(right, (dim + 1) % shape, shape) //右子节点递归创建树

        TreeNode(sorted(midIndex), dim, leftNode, rightNode)

    }
  }
}
